# -*-coding:utf-8 -*-
# @Time: 2023/3/15 18:32
# @Author: cuishuohao
# @File: demo4
# @Software: PyCharm
import numpy as np
from matplotlib import pyplot as plt

from demo1 import TrainSet, TestSet, DataSet
from demo2 import mse
from demo3 import gradient_update

w = np.random.rand()  # w和b的初始值取随机值
b = np.random.rand()
epochs = 10000  # 定义循环次数
alpha = 0.001  # 定义学习率
loss_mse = []  # 定义一个数组用于保存mse
for step in range(epochs):
    w, b = gradient_update(w, b, TrainSet, alpha)  # 更新梯度
    loss_mse.append(mse(w, b, TrainSet))
    print("训练完成，w:", w, "b:", b)

plt.scatter(TrainSet[:, 0], TrainSet[:, 1], c="blue",label="blue is train")
plt.scatter(TestSet[:, 0], TestSet[:, 1], c="red",label="red is test")
plt.plot(DataSet[:, 0], DataSet[:, 0] * w + b,label="line is predict")
plt.legend()
plt.show()

plt.plot(range(len(loss_mse)), loss_mse,label="Loss change")
plt.xlabel("Epochs")
plt.ylabel("MSE")
plt.legend()
plt.show()

print()
print("max loss is ", np.max(loss_mse))
print("mean loss is ", np.mean(loss_mse))
print("min loss is ", np.min(loss_mse))
print("finally loss is ", loss_mse[len(loss_mse)-1])
